import numpy as np
import matplotlib.pyplot as plt

T = [[3, 104, 98],
    [2, 100, 93],
    [1, 81, 95],
    [101, 10, 16],
    [99, 5, 8],
    [98, 2, 7]]

x = [18, 90]
k= 3

T = np.array(T, dtype=np.float64)
m = len(T)
x = np.array(x, dtype=np.float64)

dis = np.zeros([m, 2])
dis[:, 0] = np.sqrt((T[:, 0] - x[0])**2 + (T[:, 1] - x[1])**2)
dis[:, 1] = T[:, 2]

idx = np.argsort(dis[:, 0])
dis = dis[idx]
dis = dis[:k]
mean = dis[:, 1].mean(axis=0)
print(mean)

from sklearn.neighbors import KNeighborsRegressor
model = KNeighborsRegressor(n_neighbors=k)
model.fit(T[:, :-1], T[:, -1])
h = model.predict(x.reshape(1, -1))
print(h)
